﻿#region using
using ICSharpCode.Decompiler;
using ICSharpCode.Decompiler.Ast;
using ICSharpCode.Decompiler.ILAst;
using ICSharpCode.ILSpy;
using ICSharpCode.NRefactory;
using Mono.Cecil;
using Mono.Cecil.Cil;
using Mono.Cecil.Pdb;
using System;
using System.Collections.Generic;
using System.IO;
using System.Reflection;
using System.Text;
#endregion using

namespace PdbRestorer.Engine
{
	public class PdbGenerator
	{
		AssemblyDefinition _assembly;
		ISymbolWriter _symbolWriter;
		TextWriter _sourceCodeWriter;
		string _sourceDocumentName;
		Language _language;
		DecompilationOptions _decompOptions;
		Document _document;

		public PdbGenerator(AssemblyDefinition assembly,ISymbolWriter symbolWriter,TextWriter sourceCodeWriter,string sourceDocumentName,Language language,DecompilationOptions decompOptions)
		{
			_assembly=assembly;
			_symbolWriter=symbolWriter;
			_sourceCodeWriter=sourceCodeWriter;
			_sourceDocumentName=sourceDocumentName;
			_language=language;
			_decompOptions=decompOptions;

			DocumentLanguage docLang=language is CSharpLanguage?DocumentLanguage.CSharp:(language is ILLanguage?DocumentLanguage.Cil:DocumentLanguage.Other);
			_document=new Document(sourceDocumentName) { Language=docLang,LanguageVendor=DocumentLanguageVendor.Microsoft,Type=DocumentType.Text };
		}

		public static void Generate(string assemblyFilename,string symbolFilename,string sourceCodeFilename,string language)
		{
			if (assemblyFilename==null)
				throw new ArgumentNullException("assemblyFilename");
			if (symbolFilename==null)
				throw new ArgumentNullException("symbolFilename");
			if (sourceCodeFilename==null)
				throw new ArgumentNullException("sourceCodeFilename");
			if (language==null)
				throw new ArgumentNullException("language");

			Language lang=string.Equals(language,"C#",StringComparison.InvariantCultureIgnoreCase)||string.Equals(language,"CS",StringComparison.InvariantCultureIgnoreCase)
				 ?(Language)new CSharpLanguage()
				 :string.Equals(language,"IL",StringComparison.InvariantCultureIgnoreCase)
					  ?new ILLanguage(true)
					  :null;
			if (lang==null)
				throw new ArgumentException("Unsupported language. Use C# or IL.");
			//sourceFilename+=lang.FileExtension;

			AssemblyDefinition assemblyDefinition=AssemblyDefinition.ReadAssembly(assemblyFilename,new ReaderParameters() { AssemblyResolver=new NoErrorAssemblyResolver() });
			using (ISymbolWriter symbolWriter=new PdbWriterProvider().GetSymbolWriter(assemblyDefinition.MainModule,symbolFilename))
			using (StreamWriter sw=File.CreateText(sourceCodeFilename))
				new PdbGenerator(assemblyDefinition,symbolWriter,sw,sourceCodeFilename,lang,new DecompilationOptions())
					 .Generate();
			if (!PdbGenerator.FixSignatureAndAge(assemblyFilename,symbolFilename))
				throw new ApplicationException("Error on fixing PDB signature and age");
		}

		public void Generate()
		{
			DecompilerSettings oldDs=_decompOptions.DecompilerSettings;
			DecompilerSettings newDs;
			_decompOptions.DecompilerSettings=(newDs=_decompOptions.DecompilerSettings.Clone());
			newDs.UsingDeclarations=false;

			ExtendedPlainTextOutput output=new ExtendedPlainTextOutput(_sourceCodeWriter,ExportMemberMapping);
			DecompileAssemblyHeader(output);
			foreach (TypeDefinition td in _assembly.MainModule.GetTypes())
			{
				bool isModule=td.FullName=="<Module>";
				if (isModule)
					output.SetCommentBlock();
				_language.DecompileType(td,output,_decompOptions);
				if (isModule)
					output.ResetCommentBlock();
			}

			_decompOptions.DecompilerSettings=oldDs;
		}

		void ExportMemberMapping(MemberMapping mm)
		{
			MethodDefinition md=mm.MemberReference as MethodDefinition;
			if (md==null)
				throw new ApplicationException("Unexpected member reference");

			GenerateMethod(md,mm.LocalVariables,mm.MemberCodeMappings);
		}

		void GenerateMethod(MethodDefinition methodDefinition,IEnumerable<ILVariable> translatedLocalVariables,IEnumerable<SourceCodeMapping> sourceCodeMappings)
		{
			Mono.Cecil.Cil.MethodBody methodBody=methodDefinition.Body;

			MethodSymbols ms=new MethodSymbols(methodDefinition.MetadataToken);

			int codeSize;
			ms.GetType().GetField("code_size",BindingFlags.NonPublic|BindingFlags.Instance)
				 .SetValue(ms,codeSize=methodBody.CodeSize);
			ms.GetType().GetField("local_var_token",BindingFlags.NonPublic|BindingFlags.Instance)
				 .SetValue(ms,methodBody.LocalVarToken);
			List<int> foundVars=new List<int>();
			if (translatedLocalVariables!=null)
				foreach (ILVariable item in translatedLocalVariables)
					if ((item.OriginalVariable!=null)&&(item.OriginalVariable.Index<methodBody.Variables.Count))
					{
						VariableDefinition vd=new VariableDefinition(item.Name,item.Type);
						int index=item.OriginalVariable.Index;
						vd.GetType().GetField("index",BindingFlags.NonPublic|BindingFlags.Instance)
							 .SetValue(vd,index);
						ms.Variables.Add(vd);

						foundVars.Add(index);
					}
			foreach (VariableDefinition item in methodBody.Variables)
				if (!foundVars.Contains(item.Index))
				{
					item.Name="var"+item.Index.ToString("00");
					ms.Variables.Add(item);
				}

			SourceCodeMapping firstScm=null;
			bool offsetFound=false;
			IList<InstructionSymbol> instructionSymbols=ms.Instructions;
			foreach (SourceCodeMapping scm in sourceCodeMappings)
			{
				TextLocation startLoc=scm.StartLocation;
				TextLocation endLoc=scm.EndLocation;
				instructionSymbols.Add(new InstructionSymbol(scm.ILInstructionOffset.From,new SequencePoint(_document) { StartLine=startLoc.Line,StartColumn=startLoc.Column,EndLine=endLoc.Line,EndColumn=endLoc.Column }));

				if (scm.ILInstructionOffset.From==0)
					offsetFound=true;
				if ((firstScm==null)||(startLoc<firstScm.StartLocation))
					firstScm=scm;
			}
			if ((!offsetFound)&&(firstScm!=null))
			{
				TextLocation startLoc=firstScm.StartLocation;
				TextLocation endLoc=firstScm.EndLocation;
				instructionSymbols.Add(new InstructionSymbol(0,new SequencePoint(_document) { StartLine=startLoc.Line,StartColumn=startLoc.Column,EndLine=endLoc.Line,EndColumn=endLoc.Column }));
			}

			_symbolWriter.Write(ms);
		}

		void DecompileAssemblyHeader(ITextOutput output)
		{
			ModuleDefinition mainModule=_assembly.MainModule;

			//ICSharpCode.ILSpy.Language.DecompileAssembly
			//output.WriteLine("// "+assembly filename);
			output.WriteLine("// "+mainModule.FullyQualifiedName);
			AssemblyNameDefinition name=_assembly.Name;
			output.WriteLine("// "+(name.IsWindowsRuntime?name.Name+" [WinRT]":name.FullName));

			//ICSharpCode.ILSpy.CSharpLanguage.DecompileAssembly
			output.WriteLine();
			if (mainModule.EntryPoint!=null)
			{
				output.Write("// Entry point: ");
				output.WriteReference(mainModule.EntryPoint.DeclaringType.FullName+"."+mainModule.EntryPoint.Name,mainModule.EntryPoint,false);
				output.WriteLine();
			}
			output.WriteLine("// Architecture: "+CSharpLanguage.GetPlatformDisplayName(mainModule));
			if ((mainModule.Attributes&ModuleAttributes.ILOnly)==(ModuleAttributes)0)
				output.WriteLine("// This assembly contains unmanaged code.");
			switch (mainModule.Runtime)
			{
				case TargetRuntime.Net_1_0:
					output.WriteLine("// Runtime: .NET 1.0");
					break;
				case TargetRuntime.Net_1_1:
					output.WriteLine("// Runtime: .NET 1.1");
					break;
				case TargetRuntime.Net_2_0:
					output.WriteLine("// Runtime: .NET 2.0");
					break;
				case TargetRuntime.Net_4_0:
					output.WriteLine("// Runtime: .NET 4.0");
					break;
			}
			output.WriteLine();

			//Decompile
			DecompilationOptions options=new DecompilationOptions();
			DecompilerSettings decompilerSettings=options.DecompilerSettings;
			AstBuilder astBuilder=new AstBuilder(new DecompilerContext(_assembly.MainModule) { CancellationToken=options.CancellationToken,CurrentType=null,Settings=decompilerSettings });
			astBuilder.AddAssembly(_assembly,!options.FullDecompilation);
			astBuilder.RunTransformations();
			astBuilder.GenerateCode(output);
			output.WriteLine();
		}

		public static bool FixSignatureAndAge(string assemblyFilename,string pdbFilename)
		{
			Guid peSignature;
			int peAge;
			if (!SimplePePdbReader.GetSignatureAndAge(assemblyFilename,out peSignature,out peAge))
				return false;

			Guid pdbSignature;
			int pdbAge;
			SimplePdbFile simplePdbFile=new SimplePdbFile(pdbFilename);
			if (!simplePdbFile.Read(out pdbSignature,out pdbAge))
				return false;

			if ((pdbSignature==peSignature)&&(pdbAge==peAge))
				return true;

			return simplePdbFile.Write(peSignature,peAge);
		}

		class NoErrorAssemblyResolver:IAssemblyResolver
		{
			IAssemblyResolver _default=new DefaultAssemblyResolver();

			public AssemblyDefinition Resolve(string fullName,ReaderParameters parameters)
			{
				return _default.Resolve(fullName,parameters);
			}

			public AssemblyDefinition Resolve(string fullName)
			{
				return _default.Resolve(fullName);
			}

			public AssemblyDefinition Resolve(AssemblyNameReference name,ReaderParameters parameters)
			{
				return _default.Resolve(name,parameters);
			}

			public AssemblyDefinition Resolve(AssemblyNameReference name)
			{
				try
				{
					return _default.Resolve(name);
				}
				catch
				{
					return null;
				}
			}
		}
	}
}
